import json
import traci
import traci.constants as tc
# Assuming 'routepy' and 'logger' are imported and configured elsewhere
# Assuming VESSEL2LANE, TRUCK_INFO4RESTART, TRUCK_SUBSCRIBED_SET, TRUCK_UNSUBSCRIBED_SET are defined globally

_LANE_LENGTH_CACHE = {}
_LAST_NET_FILENAME = None

def get_cached_routepy_lane_length(net_filename, road_id):
    global _LANE_LENGTH_CACHE, _LAST_NET_FILENAME

    if net_filename != _LAST_NET_FILENAME:
        _LANE_LENGTH_CACHE.clear()
        _LAST_NET_FILENAME = net_filename

    cache_key = (net_filename, road_id)

    if cache_key in _LANE_LENGTH_CACHE:
        return _LANE_LENGTH_CACHE[cache_key]
    else:
        try:
            length = routepy.get_lane_length(net_filename, road_id)
            if length is not None:
                 _LANE_LENGTH_CACHE[cache_key] = length
            return length
        except Exception as e:
            return None


def vehicle_reset_route_stop_after_restart_simulation(current_edge, vehicle_id, configs, depart_pos, NET_FILENAME):
    # Implement the actual rerouting logic here based on your requirements
    try:
        traci.vehicle.setRoute(vehicle_id, [current_edge])
    except traci.TraCIException as e:
        pass


def vehicle_load_after_restart_simulation(configs, vessel_data_LAST, NET_FILENAME):
    global TRUCK_SUBSCRIBED_SET, TRUCK_UNSUBSCRIBED_SET,VESSEL2LANE,TRUCK_INFO4RESTART
    if vessel_data_LAST is None or len(vessel_data_LAST) ==0:
        try:
            json_vessel_data = configs.redis_conn.get("vessel_data")
            vessel_data = json.loads(json_vessel_data)
            vessel_data_LAST = vessel_data
        except Exception: # Handle potential redis/json errors silently for now
             vessel_data_LAST = {} # Fallback

    if not vessel_data_LAST: return # Cannot proceed

    vehicle_ids = list(TRUCK_INFO4RESTART.keys()) # Use list copy

    current_vessel_id_list = []
    for vessel_id in vessel_data_LAST:
        try:
            current_vessel_id_list.append(VESSEL2LANE[vessel_id]["s1_1Ts1_2"]["edge_id"])
            current_vessel_id_list.append(VESSEL2LANE[vessel_id]["s2_1Ts2_2"]["edge_id"])
        except KeyError: # Handle missing keys gracefully
            continue

    vehicle_counts = 0
    try:
         if "myType" not in traci.vehicletype.getIDList(): traci.vehicletype.add("myType", vClass="truck")
         traci.vehicletype.setMaxSpeed("myType", 50.0)
         traci.vehicletype.setAccel("myType", 5.0)
         traci.vehicletype.setDecel("myType", 9.0)
         traci.vehicletype.setTau("myType", 1.0)
    except traci.TraCIException as e_type:
        pass

    current_sim_time = traci.simulation.getTime() # Get time once

    for vehicle_id in vehicle_ids:
        if vehicle_id not in TRUCK_INFO4RESTART: continue
        truck_info = TRUCK_INFO4RESTART[vehicle_id]
        vehicle_route_edges = truck_info.get("vehicle_route_edges")
        road_id = truck_info.get("route_edges")
        depart_pos = truck_info.get("depart_pos")
        truck_stops = truck_info.get("truck_stops", [])
        depart_lane = truck_info.get("depart_lane")

        if not vehicle_route_edges or not road_id or depart_lane is None:
            continue
        route_edges_list = vehicle_route_edges.split(" ")
        if not route_edges_list: continue

        current_position_restart = None

        if depart_pos is None:
            depart_pos = 2.0

        if ":" in road_id:
            try:
                prev_edge, next_edge = routepy.find_previous_and_next_edges(route_edges_list, road_id)
                if next_edge is None: continue
                road_id_restart = next_edge
                depart_pos_restart = 1.0
            except Exception as e_find:
                continue

            try:
                current_index = route_edges_list.index(road_id_restart)
            except ValueError:
                continue

            stop_indices = []
            for stop in truck_stops:
                if not hasattr(stop, 'lane'): continue
                lane_id_stop = stop.lane
                edge_id_stop = lane_id_stop.split("_")[0]
                try:
                    stop_index = route_edges_list.index(edge_id_stop)
                    stop_indices.append(stop_index)
                except ValueError:
                    continue

            remaining_stop_edges_indices = sorted([idx for idx in stop_indices if idx >= current_index])
            remaining_stop_edges = [route_edges_list[idx] for idx in remaining_stop_edges_indices]
            temp_remaining_stops = remaining_stop_edges[:-1] if remaining_stop_edges else []
            if not any(edge.startswith("s") for edge in temp_remaining_stops):
                final_remaining_stop_edges_for_check = []
            else:
                final_remaining_stop_edges_for_check = temp_remaining_stops

            lane_length = get_cached_routepy_lane_length(NET_FILENAME, road_id_restart) # CACHE USED HERE

            if lane_length is not None:
                if depart_pos_restart >= lane_length:
                    depart_pos_restart = max(0.0, lane_length - 1.0)
                elif lane_length - depart_pos_restart < 1.0:
                    depart_pos_restart = max(0.0, depart_pos_restart - 1.1)
                current_position_restart = depart_pos_restart
            else:
                current_position_restart = depart_pos_restart

            remaining_route = route_edges_list[current_index:]
            if not remaining_route: continue

            if final_remaining_stop_edges_for_check and not any(edge in current_vessel_id_list for edge in final_remaining_stop_edges_for_check):
                route_length = len(route_edges_list)
                target_edge_reroute = route_edges_list[-1]
                if current_index in stop_indices:
                    target_index_stop = current_index + 2
                    if target_index_stop < route_length:
                        target_edge_reroute = route_edges_list[target_index_stop]
                vehicle_reset_route_stop_after_restart_simulation(target_edge_reroute, vehicle_id, configs, current_position_restart, NET_FILENAME)
                continue

            temp_route_id = f"route_{vehicle_id}_{int(current_sim_time)}_{vehicle_counts}"
            try:
                traci.route.add(temp_route_id, remaining_route)
                traci.vehicle.add(vehID=vehicle_id, routeID=temp_route_id, typeID="myType",
                                  departPos=current_position_restart, departLane="first", departSpeed=0)
                vehicle_counts+=1

                traci.vehicle.setLength(vehicle_id, 19.0)
                traci.vehicle.setColor(vehicle_id, (255, 0, 0, 255))
                traci.vehicle.setShapeClass(vehicle_id, "truck")
                traci.vehicle.setLaneChangeMode(vehicle_id, 0b010111111111)

                for stop in truck_stops:
                    if not hasattr(stop, 'lane') or not hasattr(stop, 'startPos') or not hasattr(stop, 'duration'): continue
                    lane_id_stop = stop.lane
                    lane_index_stop = int(lane_id_stop.rsplit("_", 1)[-1])
                    edge_id_stop = lane_id_stop.split("_")[0]
                    try:
                        stop_index_in_route = route_edges_list.index(edge_id_stop)
                        if stop_index_in_route >= current_index:
                            lane_pos_stop = stop.startPos
                            stop_edge_length = get_cached_routepy_lane_length(NET_FILENAME, edge_id_stop) # CACHE USED HERE

                            if stop_edge_length is not None:
                                final_stop_pos = lane_pos_stop
                                if road_id_restart == edge_id_stop and current_position_restart >= lane_pos_stop:
                                    final_stop_pos = current_position_restart + 0.1
                                final_stop_pos = max(0.0, min(final_stop_pos, stop_edge_length - 0.1))
                            else:
                                final_stop_pos = lane_pos_stop

                            traci.vehicle.setStop(vehID=vehicle_id, edgeID=edge_id_stop, laneIndex=lane_index_stop,
                                                  pos=final_stop_pos, duration=stop.duration)
                    except ValueError:
                        continue
                    except traci.TraCIException as e_stop:
                        continue
                    except Exception as e_stop_gen:
                        continue

                traci.vehicle.setSpeed(vehicle_id, -1)

            except traci.TraCIException as e_add:
                try:
                    if temp_route_id in traci.route.getIDList(): traci.route.remove(temp_route_id)
                except traci.TraCIException: pass
                continue

        elif ":" not in road_id:
            road_id_restart = road_id
            depart_pos_restart = depart_pos

            try:
                current_index = route_edges_list.index(road_id_restart)
            except ValueError:
                continue

            stop_indices = []
            for stop in truck_stops:
                if not hasattr(stop, 'lane'): continue
                lane_id_stop = stop.lane
                edge_id_stop = lane_id_stop.split("_")[0]
                try:
                    stop_index = route_edges_list.index(edge_id_stop)
                    stop_indices.append(stop_index)
                except ValueError:
                    continue

            remaining_stop_edges_indices = sorted([idx for idx in stop_indices if idx >= current_index])
            remaining_stop_edges = [route_edges_list[idx] for idx in remaining_stop_edges_indices]
            temp_remaining_stops = remaining_stop_edges[:-1] if remaining_stop_edges else []
            if not any(edge.startswith("s") for edge in temp_remaining_stops):
                final_remaining_stop_edges_for_check = []
            else:
                final_remaining_stop_edges_for_check = temp_remaining_stops

            try:
                lane_length_traci = traci.lane.getLength(depart_lane)
                if depart_pos_restart >= lane_length_traci:
                    depart_pos_restart = max(0.0, lane_length_traci - 1.0)
                elif lane_length_traci - depart_pos_restart < 1.0:
                    depart_pos_restart = max(0.0, depart_pos_restart - 1.1)
                current_position_restart = depart_pos_restart
            except traci.TraCIException as e_len:
                 current_position_restart = depart_pos_restart

            remaining_route = route_edges_list[current_index:]
            if not remaining_route: continue

            if final_remaining_stop_edges_for_check and not any(edge in current_vessel_id_list for edge in final_remaining_stop_edges_for_check):
                route_length = len(route_edges_list)
                target_edge_reroute = route_edges_list[-1]
                if current_index in stop_indices:
                    target_index_stop = current_index + 2
                    if target_index_stop < route_length:
                        target_edge_reroute = route_edges_list[target_index_stop]
                vehicle_reset_route_stop_after_restart_simulation(target_edge_reroute, vehicle_id, configs, current_position_restart, NET_FILENAME)
                continue

            temp_route_id = f"route_{vehicle_id}_{int(current_sim_time)}_{vehicle_counts}"
            try:
                traci.route.add(temp_route_id, remaining_route)
                traci.vehicle.add(vehID=vehicle_id, routeID=temp_route_id, typeID="myType",
                                  departPos=current_position_restart, departLane="first", departSpeed=0)
                vehicle_counts+=1

                traci.vehicle.setLength(vehicle_id, 19.0)
                traci.vehicle.setColor(vehicle_id, (255, 0, 0, 255))
                traci.vehicle.setShapeClass(vehicle_id, "truck")
                traci.vehicle.setLaneChangeMode(vehicle_id, 0b010111111111)

                for stop in truck_stops:
                    if not hasattr(stop, 'lane') or not hasattr(stop, 'startPos') or not hasattr(stop, 'duration'): continue
                    lane_id_stop = stop.lane
                    lane_index_stop = int(lane_id_stop.rsplit("_", 1)[-1])
                    edge_id_stop = lane_id_stop.split("_")[0]
                    try:
                        stop_index_in_route = route_edges_list.index(edge_id_stop)
                        if stop_index_in_route >= current_index:
                            lane_pos_stop = stop.startPos
                            stop_edge_length = get_cached_routepy_lane_length(NET_FILENAME, edge_id_stop) # CACHE USED HERE

                            if stop_edge_length is not None:
                                final_stop_pos = lane_pos_stop
                                if road_id_restart == edge_id_stop and current_position_restart >= lane_pos_stop:
                                    final_stop_pos = current_position_restart + 0.1
                                final_stop_pos = max(0.0, min(final_stop_pos, stop_edge_length - 0.1))
                            else:
                                final_stop_pos = lane_pos_stop

                            traci.vehicle.setStop(vehID=vehicle_id, edgeID=edge_id_stop, laneIndex=lane_index_stop,
                                                  pos=final_stop_pos, duration=stop.duration)
                    except ValueError:
                        continue
                    except traci.TraCIException as e_stop:
                        continue
                    except Exception as e_stop_gen:
                        continue

                traci.vehicle.setSpeed(vehicle_id, -1)

            except traci.TraCIException as e_add:
                try:
                    if temp_route_id in traci.route.getIDList(): traci.route.remove(temp_route_id)
                except traci.TraCIException: pass
                continue

        try:
            if vehicle_id in traci.vehicle.getIDList():
                 if vehicle_id not in TRUCK_UNSUBSCRIBED_SET:
                     if vehicle_id not in TRUCK_SUBSCRIBED_SET:
                         traci.vehicle.subscribe(vehicle_id, [
                             tc.VAR_POSITION, tc.VAR_ANGLE, tc.VAR_SPEED,
                             tc.VAR_ROAD_ID, tc.VAR_LANE_INDEX, tc.VAR_LANE_ID,
                             tc.VAR_LANEPOSITION
                         ])
                         TRUCK_SUBSCRIBED_SET.add(vehicle_id)
        except traci.TraCIException as e_sub:
            pass

    print(f"总共有{vehicle_counts}辆集卡车车被加载。")